package com.fansl.koala.quickdev.component.mybatis;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.metadata.OrderItem;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.fansl.koala.quickdev.common.exception.CheckedException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.BindingContext;
import org.springframework.web.reactive.result.method.HandlerMethodArgumentResolver;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;

/**
 * @author fansl
 * @Description: 解决Mybatis Plus Order By SQL注入问题
 * @date 2020/3/12 11:28
 */
@Slf4j
public class SqlFilterArgumentResolver implements HandlerMethodArgumentResolver {
    private final static String[] KEYWORDS = {"master", "truncate", "insert", "select"
            , "delete", "update", "declare", "alter", "drop", "sleep"};

    /**
     * 判断Controller是否包含page 参数
     *
     * @param parameter 参数
     * @return 是否过滤
     */
    @Override
    public boolean supportsParameter(MethodParameter parameter) {
        return (parameter.getParameter().getType().equals(Page.class));
    }

    /**
     * page 只支持查询 GET .如需解析POST获取请求报文体处理
     *
     * @param parameter      the method parameter
     * @param bindingContext the binding context to use
     * @param exchange       the current exchange
     * @return {@code Mono} 检查后新的page对象
     */
    @Override
    public Mono<Object> resolveArgument(MethodParameter parameter, BindingContext bindingContext, ServerWebExchange exchange) {
        MultiValueMap<String, String> queryParam = exchange.getRequest().getQueryParams();
        if (log.isDebugEnabled()) {
            log.debug("desc:{},param:{}", "page参数校验", queryParam);
        }
        String current = queryParam.getFirst("current");
        String size = queryParam.getFirst("size");
        List<String> ascs = queryParam.get("ascs");
        List<String> descs = queryParam.get("descs");
        Page page = new Page();
        if (StrUtil.isNotBlank(current)) {
            page.setCurrent(Long.parseLong(current));
        }

        if (StrUtil.isNotBlank(size)) {
            page.setSize(Long.parseLong(size));
        }
        List<OrderItem> orderItems = new ArrayList<>();
        String[] ascArray = sqlInject(ascs);
        if (ArrayUtil.isNotEmpty(ascArray)) {
            orderItems.addAll(OrderItem.ascs(ascArray));
        }
        String[] descArray = sqlInject(descs);
        if (ArrayUtil.isNotEmpty(descArray)) {
            orderItems.addAll(OrderItem.ascs(descArray));
        }
        page.setOrders(orderItems);
        if (log.isDebugEnabled()) {
            log.debug("desc:{},result:{}", "page参数校验", JSONObject.toJSON(page));
        }
        return Mono.just(page);
    }

    /**
     * SQL注入过滤
     *
     * @param strList 待验证的字符串
     */
    public static String[] sqlInject(List<String> strList) {
        String[] resultStr = new String[]{};
        if (CollectionUtil.isEmpty(strList)) {
            return resultStr;
        }
        resultStr = strList.toArray(resultStr);
        //转换成小写
        String inStr = ArrayUtil.join(resultStr, StrUtil.COMMA).toLowerCase();

        //判断是否包含非法字符
        for (String keyword : KEYWORDS) {
            if (inStr.contains(keyword)) {
                log.error("查询包含非法字符:{}", keyword);
                throw new CheckedException("包含非法字符:" + keyword);
            }
        }
        return resultStr;
    }
}
